# -*- coding: utf-8 -*-
# @Time : 2021/10/29 19:09
# @Author : Ming
# @FileName: RNN_hello_ChangeModel.py
# @Software: PyCharm


import torch
'''
如果有任何出错，会看刘二大人的官方pdf
'''
# 自定义变量
batch_size = 1      #批量数
input_size = 4      #输入张量（向量）的维度
hidden_size = 4     #中间隐层张量的维度
num_layers = 1      #RNN层数
seq_len = 5         #序列长度，也就是有多少个序列，x1,x2,x3,x4,x5

# 通过RNN实现hello转换为ohlol

idx2char = ['e', 'h', 'l', 'o']
x_data = [1, 0, 2, 2, 3]    #hello
y_data = [3, 1, 2, 3, 2]    #ohlol

# 独热编码字典
one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]]
# 对hello的每一个字母序列进行独热编码
x_one_hot = [one_hot_lookup[x] for x in x_data]

# 输入必须是 (seq_len, batch_size, input_size)
inputs = torch.Tensor(x_one_hot).view(seq_len, batch_size, input_size)
# 这里必须是长整型张量
labels = torch.LongTensor(y_data)


# Design Model
class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size, num_layers = 1):
        super(Model, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.rnn = torch.nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)

    def forward(self, input):
        hidden = torch.zeros(self.num_layers, self.batch_size, self.hidden_size)
        out, _ = self.rnn(input, hidden)
        return out.view(-1, self.hidden_size)   # Reshape out to : (seqlen*batchsize, hiddensize)


net = Model(input_size=input_size, hidden_size=hidden_size, batch_size=batch_size, num_layers=num_layers)
criterion = torch.nn.CrossEntropyLoss() #交叉熵损失
optimizer = torch.optim.Adam(net.parameters(), lr=0.05) #优化器
'''
学习率太小，也会导致效果不佳
'''

# 开始训练
for epoch in range(15):
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')
    print(', Epoch [%d/15] loss = %.3f' % (epoch + 1, loss.item()))

'''注意模型转换的对比'''
# for epoch in range(15):
#     loss = 0
#     optimizer.zero_grad()   #梯度归零
#     # 更换（小优化）模型后，不用在单独初始化h0了，已经是包含在forward里边了
#     # hidden = net.init_hidden()  # 初始化隐层h0
#     print('predicted string：', end='')
#     for input, label in zip(inputs, labels):
#         print(input.shape)
#         hidden = net(input)
#         # 这里计算loss没有使用item(),需要计算其计算图
#         loss += criterion(hidden, label)    #！！！这里hidden的预测超出了label的种类个数,更改隐层的大小即可（更改为num+1）
#         _, idx = hidden.max(dim=1)  # 找到最大下标，因为是独热，最大的1所在下标对应label
#         print(idx2char[idx.item()], end='')
#     loss.backward()
#     optimizer.step()
#     print(', epoch [%d/15] loss=%.4f' % (epoch+1, loss.item()))